/*
 * @lc app=leetcode.cn id=1382 lang=cpp
 *
 * [1382] 将二叉搜索树变平衡
 */

#include "include.h"

struct TreeNode {
    int val;
    TreeNode *left;
    TreeNode *right;
    TreeNode() : val(0), left(nullptr), right(nullptr) {}
    TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
    TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
};

// @lc code=start
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    TreeNode* balanceBST(TreeNode* root) {
        postOrderTraverse(root);

        printf("final \n");
        plotSubTree(root);
        return root;
    }

    int getBalanceFactor(TreeNode* root){
        int leftSubTreeLevel = root->left ? getHeight(root->left) : 0;
        int rightSubTreeLevel = root->right ? getHeight(root->right) : 0;
        int balanceFactor = leftSubTreeLevel - rightSubTreeLevel;
printf("get BF %d %d : %d %d \n", root->val, balanceFactor, leftSubTreeLevel, rightSubTreeLevel);
plotSubTree(root);
        return balanceFactor;
    }
    int getHeight(TreeNode* root){
        // do rebalance here when BF > 2 id found, and return the height of rebalanced tree
        int leftSubTreeLevel = root->left ? getHeight(root->left) : 0;
        int rightSubTreeLevel = root->right ? getHeight(root->right) : 0;
        return std::max(leftSubTreeLevel, rightSubTreeLevel) + 1;
    }
    void rebalanceSubTree(TreeNode* & root){
printf("before re\n");
plotSubTree(root);
        int rootBalanceFactor = getBalanceFactor(root);
        if (std::abs(rootBalanceFactor) <= 1){return;}
        // int leftBalanceFactor = root->left ? getBalanceFactor(root->left) : 0;
        // int rightBalanceFactor = root->right ? getBalanceFactor(root->right) : 0;

        switch (rootBalanceFactor)
        {
        case 2:{
            // left too heavy, need right rotate
            int leftBalanceFactor = getBalanceFactor(root->left);
            if (leftBalanceFactor == -1){
printf("2 LR\n");                
                leftRotate(root->left);
                rightRotate(root);
            }else{
printf("2 R\n");
                rightRotate(root);
            }
            break;
        }
        case -2:{
            // right too heavy, need left rotate
            int rightBalanceFactor = getBalanceFactor(root->right);
            if (rightBalanceFactor == 1){
printf("-2 RL\n");
                rightRotate(root->right);
                leftRotate(root);
            }else{
printf("-2 L\n");

show(root);
printf("ptr address %p\n", &root);
                leftRotate(root);
show(root);                
            }
            break;
        }
        default:
            break;
        }
printf("after re\n");
plotSubTree(root);
    }

    // fuck 
    void postOrderTraverse(TreeNode* & root){
        if (root->left) {postOrderTraverse(root->left);}
        if (root->right) {postOrderTraverse(root->right);}
        
        printf("%d ", root->val);
        int currentBalanceFactor = getBalanceFactor(root);
        if (std::abs(currentBalanceFactor) > 1){
            printf("re \n");
            rebalanceSubTree(root);
        }
    }
void show(TreeNode* root){
    printf("L %d %p:", root->val, root);
    if (root->left){printf("%d %p ", root->left->val, root->left);}
    else {printf("null ");}
    if (root->right){printf("%d %p ", root->right->val, root->right);}
    else {printf("null ");}
    printf("\n");
}
    void leftRotate(TreeNode* & root){
printf("ptr address %p\n", &root);
        // check error
        // if (!root->right or root->left){return;}
// plotSubTree(root);

        TreeNode* newRoot = root->right;
        root->right = newRoot->left;
        newRoot->left = root;
        root = newRoot;

        // TreeNode* temp = root;
        // root = root->right;
        // temp->right = root->left;
        // root->left = temp;

// plotSubTree(root);
    }
    void rightRotate(TreeNode* & root){
        // check error
        // if (!root->left or root->right){return;}

        // TreeNode* tempLeft;
        // tempLeft = (*root).left;
        // (*root).left = tempLeft->right;
        // tempLeft->right = &(*root);
        // root = tempLeft;

        // TreeNode* temp = root;
        // root = root->left;
        // temp->left = root->right;
        // root->right = temp;

        TreeNode* newRoot = root->left;
        root->left = newRoot->right;
        newRoot->right = root;
        root = newRoot;

        // root->right = root->left->right;
        // root->left->right = root;
        // root = root->left;
    }

    // for test
    void plotSubTree(TreeNode* root){
        // int height = getHeight(root);
        // int fullBinaryTreeSize = std::pow(2, height);
        // std::vector<int> binaryHeap(fullBinaryTreeSize, 0);
        ;
        levelOrderTraverse(root);
    }
    void levelOrderTraverse(TreeNode* root) {
        if (!root) {return;}
        std::queue<TreeNode*> nodePtrQueue;
        nodePtrQueue.push(root);
        while (!nodePtrQueue.empty())
        {
            TreeNode* topNodePtr = nodePtrQueue.front();
            nodePtrQueue.pop();

            printf("%d ", topNodePtr->val);

            if (topNodePtr->left) {nodePtrQueue.push(topNodePtr->left);}
            else {printf("null ");}
            if (topNodePtr->right) {nodePtrQueue.push(topNodePtr->right);}
            else {printf("null ");}
        }
        // for test
        printf("\n");
    }
};
// @lc code=end

int main(){
    TreeNode* root = new TreeNode;
    root->val = 1;
    TreeNode* node2 = new TreeNode;
    node2->val = 2;
    TreeNode* node3 = new TreeNode;
    node3->val = 3;
    TreeNode* node4 = new TreeNode;
    node4->val = 4;

    root->right = node2;
    node2->right = node3;
    node3->right = node4;

    Solution test;
    test.plotSubTree(root);

    test.balanceBST(root);

    delete root;
    delete node2;
    delete node3;
    delete node4;
}

// [19,10,null,4,17,null,5]
